"""--------------------------------------------------------------------------------------------------
Filename:   database.py
Author:     Ben Han
Date:       20-Apr-2010
Description:
    Contains classes to process and store data for the biofeedback application.
--------------------------------------------------------------------------------------------------"""

from __future__ import with_statement   # to use the automatic file.close() functionality
import numpy as np
import wx, os, sys, traceback
import time
from config import ConfigData
import pickle


#===============================================================================
# CLASS: Patient
#    This class stores data for each patient. 
#    It is the top level object which contains all other data objects.
#    Must be stored as a native python object. The swig wrappers for the C++ 
#    wx code don't like to be pickled.
#===============================================================================
class Patient:
    def __init__(self, id=False, name='', dob='', gender='', notes='', sessions=[]):
        # Initial attributes
        if not id: 
            id = str(int(time.time()))  # Create an arbitrary unique ID if none provided.
            id = id[-8:]                # Uses the number of seconds since 00:00:00 1/1/1970
        self.id = id                        # patient ID number    
        self.name = name                    # patient name
        self.dob = dob                      # patient date of birth
        self.gender = gender                # patient gender
        self.notes = notes                  # storage for general notes, condition, treatment, etc
        self.session_record = sessions      # array containing references to the sessions completed by this patient
        temp = wx.DateTime.Now()
        self.date = temp.Format()           # date created, set to current time & date.
            # Format() returns a string in a format that can be parsed
            # back into a wx.DateTime object after unpickling.
        self.path = ""
        self.cfg = ConfigData()
    
    def __str__(self):
        s = 'ID:\t'+ str(self.id)+ '\n'+\
            'Name:\t'+ self.name+ '\n'+\
            'DoB:\t'+ self.dob+ '\n'+\
            'Gender:\t'+ self.gender+ '\n'+\
            'Date:\t'+ self.date+ '\n'+\
            'No. of sessions: '+ str(len(self.session_record))+ '\n'+\
            'Notes:\n'+ self.notes
        return s
        
    def add_session(self, session_rec):
        # Add a session to the internal record
        self.session_record.append(session_rec)
        return len(self.session_record)
        
    def export(self, parent):
        # Opens up a directory selection dialog for user to choose location for output
        # CSV files.
        patdir = self.path[:self.path.rfind('\\') + 1]
        dlg = wx.DirDialog(parent,
                           message="Choose a location for output files",
                           defaultPath=patdir)

        path = ""
        if dlg.ShowModal() == wx.ID_OK:
            path = dlg.GetPath()    # directory name for output
            patfilename = self.path[self.path.rfind('\\')+1:self.path.rfind('.pat')]
            outdir = path + "\\" + patfilename + "_output"
            # Create a folder for output
            try:
                os.makedirs(outdir)
            except OSError:
                pass    # directory already exists
            except:
                sys.stderr.write(traceback.format_exc())
            else:
                path = outdir
            
            for i in range(len(self.session_record)):
                sr = self.session_record[i]
                j = self.path.rfind('\\')
                spath = self.path[:j+1] + sr.path
                try:
                    file = open(spath, 'rb')
                    file.flush()
                    s = pickle.load(file)
                    file.close()
                except:
                    msg = ">> Error loading session from file.\n"
                    sys.stderr.write(msg)
                    wx.MessageDialog(None, msg, "Session load error", wx.ICON_WARNING | wx.OK).ShowModal()
                    sys.stderr.write(traceback.format_exc())
                else:
                    outpath = path+ "\\"+ patfilename+ "_Sess"+ str(i+1).rjust(3,'0')
                    s.export(parent, outpath)
        return path
        


#===============================================================================
# CLASS: SessionRecord
#    Database object used to represent session objects in the patient object
#    Actual session objects should be saved in the same folder as the patient
#===============================================================================
class SessionRecord:
    def __init__(self, s_obj):
        # Constructor takes a database.Session object as its single argument
        self.date = s_obj.date
        self.type = s_obj.type
        self.stats = s_obj.stats
        self.path = ""    # path to the file containing the pickled session object
        

#===============================================================================
#  CLASS: Session
#     Stores the session information and configuration file used 
#===============================================================================
class Session:
    def __init__(self, config, pid, type, blocks, trials, stats):
        self.timestamp = str(int(time.time()))
        temp = wx.DateTime.Now()
        self.date = temp.Format()           # time and date of session creation
                            # Date must be stored as a string here, as in Patient class (to be pickled)
        self.config = config                # ConfigData object used for the session
        self.patient = pid                  # ID of patient
        self.type = type                    # type of training (string)
        self.blocks = blocks                # list of block objects for the session
        self.trials = trials                # list of trial objects for the session
        self.stats = stats                  # session statistics
        self.calib_mean = 0                 # mean value of calibration peaks
        
        # Initialise stats object
        self.calc()
        
    def __str__(self):
        s = 'Patient ID,'+ str(self.patient)+ '\n'+\
            'Date,'+ self.date+ '\n'+\
            'Training,'+ self.type+ '\n'+\
            'STATISTICS,'+ '\n'+\
            self.stats.get_col_labels()+ str(self.stats)
        return s
        
    def add_block(self, block):
        self.blocks.append(block)
        
    def add_trial(self, trial):
        self.trials.append(trial)
        self.calc()
        
    def calc(self):
        # Calculate session statistics.
        # The results are written directly to the local Stats object (stats)
        if not self.stats:
            self.stats = Stats()
        s = self.stats
        s.nhits = 0
        s.ntrials = len(self.trials)
        if s.ntrials == 0:
            pass
        else:
            peaks = []
            for t in self.trials:
                peaks.append(t.peak[1])     # extract the EMG value (float64)
                if t.result:
                    s.nhits+= 1
                
            data = np.array(peaks)
            s.min = data.min()
            s.max = data.max()
            s.mean = data.mean()
            s.hitrate = (s.nhits * 100) / s.ntrials
        
    def export(self, parent, filename=False):
        """ Write session to CSV files. Two files will be created: 
            1- Store stats for the session and each block. 
                It will also record the history of hits and misses. 
            2- The second file will store all the raw EMG data gathered 
                during the session.
            """
        if not filename:
            filename = 'ID_'+ str(self.patient)+ '_'+ self.timestamp
            file_choices = "Comma Seperated Values (*.csv)|*.csv"
            dlg = wx.FileDialog(
                parent,
                message="Export session data to CSV files",
                defaultDir=os.getcwd(),
                defaultFile=filename,
                wildcard=file_choices,
                style=wx.SAVE)
    
            if dlg.ShowModal() == wx.ID_OK:
                filename = dlg.GetPath()
            else:
                filename = ""       # cancelled
                return filename
                
        if filename[-4:] == '.csv':
            filename = filename[:-4]
        self.write_stats(filename)
        self.write_emg(filename)
        return filename
    
    def write_stats(self, filename):
        # write session statistics to a CSV file 
        name = filename+ '_stats.csv'
        file = open(name, 'w')
        file.write('<Output for session statistics>\n')
        file.write('Patient ID,'+ str(self.patient)+ '\n')
        file.write('Session date,'+ str(self.date)+ '\n')
        file.write('Training type,'+ str(self.type)+ '\n')
        file.write('PROTOCOLS:\n')
        file.write('Delay,%s,length,%s\n' % (str(self.config._use_delay), str(self.config._delay_len)))
        file.write('Fading,%s\n' % str(self.config._use_fading))
        file.write('Calibration mean peak value,%s\n\n' % str(self.calib_mean))
        file.write('Block Stats\n')
        file.write(self.stats.get_col_labels())
        bi = 1
        for b in self.blocks:       # for each block in the blocks list...
            # write block statistics to file 
            file.write(str(bi)+ ','+ str(b.stats))
            bi += 1
        file.write("\n")
        file.write('Session total,'+ str(self.stats))     # write total session stats
        file.write("\n")
        file.write('Summary of trials\n')
        if len(self.trials)>0:
            file.write(self.trials[0].get_col_labels())
            ti = 1
            for tr in self.trials:
                # write CSV string representation of the history to file
                file.write(str(ti)+ ','+ str(tr))
                ti += 1
        else:
            file.write('(no record of trials)')
        file.flush()      # flush file buffer
        os.fsync(file)
        file.close()
        
    def write_emg(self, filename):
        # write EMG data to a CSV file
        fname = filename+ '_emg.csv'
        file = open(fname, 'w')
        file.write('<Output for session data (raw EMG values)>\n')
        file.write('Patient ID,'+ str(self.patient)+ '\n')
        file.write('Session date,'+ str(self.date)+ '\n')
        file.write('Training type,'+ str(self.type)+ '\n')
        file.write('\n')
        
        # Create a numpy array (2-D) containing all the emg data from trials
        # Every column represents one trial
        ntrials = len(self.trials)         # number of trials in the list
        nrow = self.config._trial_len*self.config._sample_rate + 1   # this needs to be able to handle incomplete trials
        sarray = np.zeros( (nrow, 2*ntrials) )   # initialise a NumPy array filled with zeroes
        
        for c in range(ntrials):               # loop over each column of the array (one trial per column)
            emg = self.trials[c].emgdata    # grab a reference to the emgdata list of each trial
            for r in range(len(emg)):       # loop over time values
                sarray[r,2*c] = emg[r][1]       # place emg value in the array
                if emg[r][2]:
                    sarray[r,2*c+1] = 1     # place plot state in the array
                else:
                    sarray[r,2*c+1] = 0
        # Write column labels to file stream
        labels = 'Time'
        for c in range(ntrials):
            labels+= ','+ 'Trial '+ str(c+1)+ ',Plot visible'
        labels+= '\n'
        file.write(labels)
        
        # Loop over the array and write EMG values to file
        # Each row represents the EMG value for the same time co-ordinate from each trial.
        # The result is a CSV output with one trial per column.
        from decimal import Decimal
        sample_period = Decimal(1) / Decimal(str(self.config._sample_rate))
        for r in range(nrow):       # loop over number of rows
            entry = str(r*sample_period) + ','           # time value
            for c in range(ntrials):   # loop over number of columns
                entry+= "%.2f,%s," % ( float(sarray[r,2*c]), str(sarray[r,2*c+1]) )    # format EMG value
            entry+= '\n'
            file.write(entry)
            
        file.flush()    # flush file buffer
        os.fsync(file)  # ensure the file buffers are written to disk
        file.close()


#===============================================================================
#  CLASS: Block
#     Stores block information and raw EMG data 
#===============================================================================
class Block:
    def __init__(self, history=[], stats=False):
        self.history = history              # record of peaks, targets, hits
        self.stats = stats                  # block statistics
        
    def calc(self):
        # Calculate block stats
        #    <history> list format: 
        #        ( (float X, float Y) PEAK, string TARGET_DETAILS, bool HIT)
        if not self.stats:
            self.stats = Stats()
        s = self.stats
        s.nhits = 0
        s.ntrials = len(self.history)
        if s.ntrials == 0:
            pass
        else:
            y = []
            for i in self.history:
                y.append(i[0][1])
                if i[2]: s.nhits+= 1
            data = np.array(y)
            s.min = data.min()
            s.max = data.max()
            s.mean = data.mean()
            s.hitrate = (s.nhits * 100) / s.ntrials


#===============================================================================
# CLASS: Trial
#    Class for storing trial details
#===============================================================================
class Trial:
    def __init__(self, len=None, typ=None, peak=None, mpl_obj=None, result=None, emgdata=[], plot_size=(), y_max=0.0, aspect=True):
        self.length = len           # length of trial in seconds
        self.type = typ             # training type (string)
        self.peak = peak            # a tuple containing the peak coordinates
        self.result = result        # set to True if target was hit (boolean)
        self.emgdata = emgdata      # list of tuples storing the raw emg data collected (time, value)
        self.target = self.get_target_tuple(mpl_obj)    # tuple storing parameters of matplotlib Rectangle, or Line2D
        self.plot_size = plot_size  # tuple of the plot size in pixels (width, height)
        self.y_max = y_max          # float value of y axis range (in uV)
        self.aspect = aspect
        
    def __str__(self):
        # Returns a CSV formatted summary of this trial
        len = "%i," % int(self.length)
        # get the string representation for the peak
        peak = "%.2f,%.2f," % ( float(self.peak[0]), float(self.peak[1]) )
        # get the string representation for the target
        targ = self.get_target_str()
        # get the string representation for the result
        if self.result:    # if target was hit...
            res = '1,'
        else:
            res = '0,'
        # get the string representation for the plot size and height
        plot = "%i,%i,%.2f," % (self.plot_size[0], self.plot_size[1], self.y_max)
        # get the actual delay (time plot shown - time for peak value)
        
        # add it all together and return
        s = len+ res+ peak+ plot+ targ+ '\n'
        return s
            
    def get_target_tuple(self, mpl_obj=0):
        # convert the target matplotlib object into a primitive python tuple
        if mpl_obj == 0:
            target = self.target
        else:
            target = mpl_obj
            
        from matplotlib.patches import Rectangle
        from matplotlib.lines import Line2D
        rect = Rectangle( (0,0), 0, 0 )
        line = Line2D( (0,1), (0,1) )
        if type(target) == type(rect):
            # for skill (x, y, width, height) where x and y refer to the coordinate of the lower left corner
            return (target.get_x(), target.get_y(), target.get_width(), target.get_height())
        elif type(target) == type(line):
            return tuple(target.get_ydata()) 
        else:
            return ()
        
    def get_target_str(self, output=0):
        # Converts the target tuple into a string
        #    Pass in output = 0 for CSV formatted output
        #    Pass in output = 1 for list format (used in trial details of session pane)
        if not self.target:
            return '(none)'
        
        r = self.target
        if len(r) == 2:
            ts = "%.2f" % float(r[0])
            return ts
        elif len(r) == 4:
            cx = r[0]+ r[2]/2.0     # calculate the centre coordinates
            cy = r[1]+ r[3]/2.0     # calculate the centre coordinates
            x = "%.2f" % float(cx)          # centre X
            y = "%.2f" % float(cy)          # centre Y
            w = "%.2f" % float(r[2])        # width
            h = "%.2f" % float(r[3])        # height
            a = "%.2f" % float(r[2]*r[3])   # area
            f = str(self.aspect)            # fixed aspect ratio, yes or no
            pd = "%.2f" % (float(r[0]) + float(r[2]) - float(self.peak[0]))
            
            if output == 0:
                # Convert the matplotlib.patches.Rectangle argument into a string in CSV format
                s = w+ ','+ h+ ','+ a+ ','+  x+ ','+ y+ ','+ pd+ ','+ f+ ','
                return s
            elif output == 1:
                # Convert rectangle to human readable format
                s = 'W:%s, H:%s, Area:%s, Centre:(%s, %s), Fixed aspect ratio: %s' % (w, h, a, x, y, f)
                return s
            else:
                return 'invalid output type'
        else:
            return 'invalid target'

    def emg_str(self):
        # Returns CSV formatted string of the emg data for this trial
        s = ''
        for e in self.emgdata:
            s+= ("%.2f" % float(e[0])+ ','+ "%.2f" % float(e[1])+ '\n')
        return s
        
    def get_col_labels(self):
        if len(self.target) == 2:
            s = 'Trial Number,Length(secs),Result(1=hit),Peak X (secs),Peak Y (uV),Plot size X (pixels),Plot size Y (pixels),Y axis max (uV),Target amplitude\n'
        elif len(self.target) == 4:
            s = 'Trial Number,Length(secs),Result(1=hit),Peak X (secs),Peak Y (uV),Plot size X (pixels),Plot size Y (pixels),Y axis max (uV),Target W (secs),Target H (uV),Target Area,Centre X (secs),Centre Y (uV),Peak delay,Fixed Aspect Ratio\n'
        return s
    
    def get_summary(self):
        # Prints a summary of the trial parameter to sys.stdout
        trial_s = '> Trial summary\n'
        trial_s += 'Training mode: %s\n' % str(self.type)
        trial_s += 'Length: %.2f seconds\n' % float(self.length)
        trial_s += 'Plot size: (%i, %i) pixels\n' % (self.plot_size[0], self.plot_size[1])
        trial_s += 'Y axis max: %.2f uV\n' % self.y_max
        trial_s += 'Target: %s\n' % self.get_target_str(1)
        trial_s += 'Peak: (%.2f, %.2f)\n' % (float(self.peak[0]), float(self.peak[1]))
        trial_s += 'Result: %s\n' % str(self.result)
        return trial_s


#===============================================================================
#  CLASS: Stats
#     Used by Block and Session to store basic stats 
#===============================================================================
class Stats:
    def __init__(self):
        self.ntrials = 0
        self.min = 0
        self.max = 0
        self.mean = 0
        self.nhits = 0
        self.hitrate = 0
        
    def __str__(self):
        # The return value is a string containing the class attributes in CSV format.
        s = "%i" % int(self.ntrials)+ ','+ "%.2f" % float(self.min)+ ','+ "%.2f" % float(self.max)+\
            ','+ "%.2f" % float(self.mean)+ ','+ "%i" % int(self.nhits)+ ','+ "%.2f" % float(self.hitrate)+ '\n'
        return s
    
    def __repr__(self):
        # The representation entered into the patient database
        s = ''
        s+= 'No. of trials: '+ "%i" % int(self.ntrials)+ '\n'
        s+= 'Min peak:      '+ "%.2f" % float(self.min)+ '\n'
        s+= 'Max peak:      '+ "%.2f" % float(self.max)+ '\n'
        s+= 'Mean peak:     '+ "%.2f" % float(self.mean)+ '\n'
        s+= 'No. of hits:   '+ "%i" % int(self.nhits)+ '\n'
        s+= 'Hit rate:      '+ "%.2f" % float(self.hitrate)+ '\n'
        return s
    
    def get_col_labels(self):
        s = 'Block No.,No. trials,Min peak,Max peak,Mean peak,No. hits,Hitrate\n'
        return s

# End of database.py
